import torch
from torch import nn


class Encoder(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        """
        编码器的目的：传递隐藏状态给Decoder
        :param x: 编码器输入序列
        :return: 隐藏状态
        """
        raise ImportError()


class Decoder(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x, h):
        """
        解码器的目的，接收编码器Encoder传递的隐藏状态
        :param x: 解码器的输入序列
        :param h: Encoder编码器的隐藏状态
        :return: 解码器的输出序列
        """


class EncoderDecoder(nn.Module):
    def __init__(self, encoder, decoder):
        super().__init__()

        self.encoder = encoder
        self.decoder = decoder

    def forward(self, x1, x2):
        """
        融合编解码器
        :param x1: 编码器Encoder的输入
        :param x2: 解码器Decoder的输入(某种程度上也是labels输出)
        :return: 解码器输出序列
        """
        h = self.encoder(x1)
        y = self.decoder(x2, h)
        return y
